import pandas as pd


def load_titanic_data(train_path='train.csv', test_path='test.csv'):
    """加载泰坦尼克号数据集"""
    try:
        train_data = pd.read_csv(train_path)
        test_data = pd.read_csv(test_path)
        print(f"训练数据: {train_data.shape}, 测试数据: {test_data.shape}")
        return train_data, test_data
    except Exception as e:
        print(f"数据加载错误: {e}")
        return None, None


def preprocess_data(df):
    # 复制数据以避免链式赋值警告
    df = df.copy()

    # 填充缺失值
    df['Age'] = df['Age'].fillna(df['Age'].median())
    df['Fare'] = df['Fare'].fillna(df['Fare'].mean())
    df['Embarked'] = df['Embarked'].fillna(df['Embarked'].mode()[0])

    # 添加 FamilySize 列
    df['FamilySize'] = df['SibSp'] + df['Parch'] + 1

    # 添加 IsAlone 列
    df['IsAlone'] = (df['FamilySize'] == 1).astype(int)

    # 修复：使用原始字符串修复正则表达式警告
    df['Title'] = df['Name'].str.extract(r' ([A-Za-z]+)\.', expand=False)
    df['Title'] = df['Title'].replace(
        ['Lady', 'Countess', 'Capt', 'Col', 'Don', 'Dr', 'Major', 'Rev', 'Sir', 'Jonkheer', 'Dona'], 'Rare')
    df['Title'] = df['Title'].replace('Mlle', 'Miss')
    df['Title'] = df['Title'].replace('Ms', 'Miss')
    df['Title'] = df['Title'].replace('Mme', 'Mrs')

    # 将分类变量转换为数值
    df['Sex'] = df['Sex'].map({'male': 0, 'female': 1})
    df['Embarked'] = df['Embarked'].map({'S': 0, 'C': 1, 'Q': 2})
    df['Title'] = df['Title'].map({"Mr": 1, "Miss": 2, "Mrs": 3, "Master": 4, "Rare": 5})

    # 删除不必要的列（保留PassengerId）
    df = df.drop(['Ticket', 'Cabin', 'Name'], axis=1)

    return df